package org.junyee.demo.security.ec;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.math.BigInteger;
import java.security.spec.AlgorithmParameterSpec;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

import org.junyee.demo.security.exception.KeyException;
import org.springframework.util.Assert;
import org.springframework.util.Base64Utils;
import org.springframework.util.ObjectUtils;
import org.springframework.util.StringUtils;
import sun.security.util.*;
import sun.security.x509.*;
import lombok.AllArgsConstructor;
import lombok.Getter;
import lombok.extern.slf4j.Slf4j;

@Slf4j
@SuppressWarnings("restriction")
public class PemUtils {

    private static final String FORMAT_PUBLIC_PEM_HEAD = "-----BEGIN PUBLIC KEY-----";
    private static final String FORMAT_PUBLIC_PEM_FOOT = "-----END PUBLIC KEY-----";

    private static final String FORMAT_PRIVATE_PEM_HEAD = "-----BEGIN PRIVATE KEY-----";
    private static final String FORMAT_PRIVATE_PEM_FOOT = "-----END PRIVATE KEY-----";

    public static String enhance(String encodeText) {
        Assert.hasText(encodeText, "encode text not be empty");
        StringBuilder text = new StringBuilder();
        int len = encodeText.length();
        for (int i = 0; i < len;) {
            int end = i + 64;
            if (i > 0) {
                text.append("\r\n");
            }
            text.append(encodeText.substring(i, i = end > len ? len : end));
        }
        return text.toString();
    }

    public static boolean isPkcs8PublicKeyFormat(String pem) {
        Assert.hasText(pem, "PEM Requried");
        pem = StringUtils.trimWhitespace(pem);
        return pem.startsWith(FORMAT_PUBLIC_PEM_HEAD) && pem.endsWith(FORMAT_PUBLIC_PEM_FOOT);

    }

    public static boolean isPkcs8PrivateKeyFormat(String pem) {
        Assert.hasText(pem, "PEM Requried");
        pem = StringUtils.trimWhitespace(pem);
        return pem.startsWith(FORMAT_PRIVATE_PEM_HEAD) && pem.endsWith(FORMAT_PRIVATE_PEM_FOOT);
    }

    public static byte[] getBytes(String pem) {
        Assert.hasText(pem, "PEM Requried");
        pem = StringUtils.trimWhitespace(pem);
        boolean isPublicKey = isPkcs8PublicKeyFormat(pem);
        boolean isPrivatekey = isPkcs8PrivateKeyFormat(pem);
        if (!isPublicKey && !isPrivatekey) {
            throw new KeyException("Not pkcs8 format PEM");
        }
        int start = isPublicKey ? 26 : 27;
        int end = isPublicKey ? 24 : 25;
        pem = StringUtils.trimAllWhitespace(pem.substring(start, pem.length() - end));
        return Base64Utils.decodeFromString(pem);
    }

    public static PemEncodedKeySpec getKeySpec(String pem) {
        byte[] bytes = getBytes(pem);
        boolean isPublicKey = isPkcs8PublicKeyFormat(pem);
        X509EncodedKeySpec x509EncodedKeySpec = isPublicKey ? new X509EncodedKeySpec(bytes)
                : getPublicKeyFromPrivateKey(bytes);
        PKCS8EncodedKeySpec pkcs8EncodedKeySpec = isPublicKey ? null : new PKCS8EncodedKeySpec(bytes);
        return new PemEncodedKeySpec(x509EncodedKeySpec, pkcs8EncodedKeySpec);
    }

    /**
     * Get public key in the private Key if contain a public key. It will be
     * resolved by ASN.1 PKCS8 Format
     * 
     * @param bytes The private key bytes
     * @return
     */
    public static X509EncodedKeySpec getPublicKeyFromPrivateKey(byte[] bytes) {
        Assert.isTrue(!ObjectUtils.isEmpty(bytes), "encodeBytes must be not empty");
        try {
            // get private key
            DerValue derValue = new DerValue(new ByteArrayInputStream(bytes));
            // ASN.1 PKCS8 defined start with sequence
            if (derValue.tag != DerValue.tag_Sequence) {
                throw new KeyException("invalid key format");
            }
            // ASN.1 pkcs8 defined version is 0.
            BigInteger version = derValue.data.getBigInteger();
            if (!version.equals(BigInteger.ZERO)) {
                throw new KeyException("pkcs8 version mismatch: (supported:0)");
            }
            // ASN.1 pkcs8 defined OID and built it to be algorithmId
            AlgorithmId algid = AlgorithmId.parse(derValue.data.getDerValue());
            // ASN.1 pkcs8 defined ECC Private Key bytes.
            bytes = derValue.data.getOctetString();
            derValue = new DerValue(new ByteArrayInputStream(bytes));
            // ASN.1 pkcs8 defined ECC Private Key start with sequence
            if (derValue.tag != DerValue.tag_Sequence) {
                throw new KeyException("invalid key format");
            }
            // ASN.1 pkcs8 Sec1 defined ECC Private Key version 1
            int version1 = derValue.data.getInteger();
            if (version1 != 1) {
                throw new KeyException("ecc private key version mismatch: (supported:1)");
            }
            // ASN.1 pkcs8 Sec1 defined ECC Private Key property d.
            bytes = derValue.data.getOctetString();
            // ASN.1 pkcs8 Sec1 defined ecc optional parameters if contain some optional
            if (derValue.data.available() != 0) {
                DerOutputStream out = new DerOutputStream();
                DerValue derValue3 = derValue.data.getDerValue();
                // 0xA0 : ASN.1 pkcs8 [0] optional which is OID
                if (derValue3.tag == (byte) 0xA0) {
                    ObjectIdentifier oid = derValue3.data.getOID();
                    if (log.isDebugEnabled()) {
                        log.debug("[0] OID: {} ", oid);
                    }
                    DerValue params = null;
                    if (derValue3.data.available() == 0) {
                        params = null;
                    } else {
                        params = derValue3.data.getDerValue();
                        if (params.tag == DerValue.tag_Null) {
                            if (params.length() != 0) {
                                throw new IOException("invalid NULL");
                            }
                            params = null;
                        }
                        if (derValue3.data.available() != 0) {
                            throw new IOException("Invalid AlgorithmIdentifier: extra data");
                        }
                    }
                    derValue3 = derValue.data.available() != 0 ? derValue.data.getDerValue() : null;
                }
                // 0xA1 : ASN.1 pkcs8 Sec1 [1] optional which is public key
                if (derValue3 != null && derValue3.tag == (byte) 0xA1) {
                    // convert to ASN.1 pkcs8 Sec1 defined public key format
                    encodePublicKey(out, algid, derValue3.data.getUnalignedBitString());
                    return new X509EncodedKeySpec(out.toByteArray());
                }

            }
        } catch (IOException e) {
            throw new KeyException(e);
        }
        if (log.isWarnEnabled()) {
            log.warn("Can not found public key!");
        }
        return null;
    }

    private static void encodePublicKey(DerOutputStream out, AlgorithmId algid, BitArray key) {
        DerOutputStream tmp = new DerOutputStream();
        try {
            algid.encode(tmp);
            tmp.putUnalignedBitString(key);
            out.write(DerValue.tag_Sequence, tmp);
        } catch (IOException e) {
            throw new KeyException(e);
        }

    }

    @AllArgsConstructor
    @Getter
    public static class PemEncodedKeySpec implements AlgorithmParameterSpec {
        private X509EncodedKeySpec publicKey;
        private PKCS8EncodedKeySpec privateKey;
    }

}
